import json
import os
from PIL import Image, ImageDraw, ImageFont
from tqdm import tqdm
import time

# --- 配置区 ---
# 颜色
CORRECT_COLOR = (0, 180, 0)       # 柔和的绿色
INCORRECT_COLOR = (255, 50, 50)   # 柔和的红色
QUERY_COLOR = (0, 100, 255)       # 蓝色
TEXT_BG_COLOR = (255, 255, 255, 200) # 半透明白色背景 (透明度略微调高)
SEPARATOR_COLOR = (255, 0, 0)     # 分割线颜色 (红色)

# 尺寸
SEPARATOR_WIDTH = 5               # 分割线宽度 (像素)

# 字体设置 (尝试加载常用字体，如果失败则使用Pillow默认字体)
try:
    FONT = ImageFont.truetype("arial.ttf", 20) # 稍微调小字体以适应更多信息
except IOError:
    print("警告: Arial 字体未找到，将使用默认字体。")
    FONT = ImageFont.load_default()

def get_item_id_from_filename(filename):
    """从 '1213_4564.jpg' 格式的文件名中提取货号ID '1213'"""
    try:
        return filename.split('_')[0]
    except IndexError:
        return os.path.splitext(filename)[0]

def create_placeholder_image(size=(300, 300), text="Image\nNot Found"):
    """当图片文件不存在时，创建一个占位图"""
    img = Image.new('RGB', size, color='gray')
    draw = ImageDraw.Draw(img)
    try:
        draw.text((size[0]/2, size[1]/2), text, fill='white', font=FONT, anchor='mm', align='center')
    except TypeError:
        text_width, text_height = draw.textsize(text, font=FONT)
        position = ((size[0] - text_width) / 2, (size[1] - text_height) / 2)
        draw.text(position, text, fill='white', font=FONT)
    return img

def draw_text_with_dynamic_bg(draw, position, text, font, text_color, bg_color):
    """
    在图片上绘制带有动态大小背景的文本。
    """
    try:
        # 使用 textbbox 计算精确的边界框 (Pillow 8.0.0+)
        text_bbox = draw.textbbox(position, text, font=font, spacing=4)
        # 增加一点内边距 (padding)
        bg_bbox = [text_bbox[0] - 5, text_bbox[1] - 5, text_bbox[2] + 5, text_bbox[3] + 5]
        draw.rectangle(bg_bbox, fill=bg_color)
        draw.text(position, text, font=font, fill=text_color, spacing=4)
    except AttributeError:
        # 兼容旧版 Pillow
        draw.text(position, text, font=font, fill=text_color, spacing=4)
        print("警告: 您的Pillow版本较旧，文本背景可能不精确。建议升级 (pip install --upgrade Pillow)。")


def visualize_one_result(
    query_image_filename,
    similarity_data_for_query,
    image_dir,
    output_path,
    top_n=10,
    std_height=300
):
    """
    为单个查询图片生成并保存可视化结果。这是一个被主循环调用的辅助函数。
    """
    query_item_id = get_item_id_from_filename(query_image_filename)
    similar_images_data = similarity_data_for_query[:top_n]
    
    processed_images = []

    # 1. 处理查询图片
    query_image_path = None
    for img_dir in image_dir:
        path_candidate = os.path.join(img_dir, query_image_filename + '.jpg')
        if os.path.isfile(path_candidate):
            query_image_path = path_candidate
            break
    
    try:
        if query_image_path is None: raise FileNotFoundError
        img = Image.open(query_image_path).convert("RGBA")
    except FileNotFoundError:
        img = create_placeholder_image(text=f"{query_image_filename}\nNot Found")
    
    aspect_ratio = img.width / img.height
    new_width = int(std_height * aspect_ratio)
    img = img.resize((new_width, std_height), Image.LANCZOS)
    
    draw = ImageDraw.Draw(img)
    # --- 修改点 2: 为Query图也加上文件名 ---
    query_text = f"Query\n{query_image_filename}"
    # --- 修改点 1: 使用函数绘制动态背景的文本 ---
    draw_text_with_dynamic_bg(draw, (10, 8), query_text, FONT, QUERY_COLOR, TEXT_BG_COLOR)
    processed_images.append(img)
    
    # 2. 处理相似图片
    for result in similar_images_data:
        similar_filename, score = list(result.items())[0]
        
        similar_image_path = None
        for img_dir in image_dir:
            path_candidate = os.path.join(img_dir, similar_filename + '.jpg')
            if os.path.isfile(path_candidate):
                similar_image_path = path_candidate
                break

        try:
            if similar_image_path is None: raise FileNotFoundError
            img_sim = Image.open(similar_image_path).convert("RGBA")
        except FileNotFoundError:
            img_sim = create_placeholder_image(text=f"{similar_filename}\nNot Found")
            
        aspect_ratio_sim = img_sim.width / img_sim.height
        new_width_sim = int(std_height * aspect_ratio_sim)
        img_sim = img_sim.resize((new_width_sim, std_height), Image.LANCZOS)

        similar_item_id = get_item_id_from_filename(similar_filename)
        is_correct = (query_item_id == similar_item_id)
        
        status_text = "Correct" if is_correct else "Incorrect"
        text_color = CORRECT_COLOR if is_correct else INCORRECT_COLOR

        draw_sim = ImageDraw.Draw(img_sim)
        text_to_draw = f"{similar_filename}\nSim: {score:.3f}\n{status_text}"
        
        # --- 修改点 1: 使用函数绘制动态背景的文本 ---
        draw_text_with_dynamic_bg(draw_sim, (10, 8), text_to_draw, FONT, text_color, TEXT_BG_COLOR)
        
        processed_images.append(img_sim)

    # --- 修改点 3: 拼接图片并添加红色分割线 ---
    if not processed_images: return # 如果没有图片，则不处理

    # 计算总宽度，包括所有图片宽度和分割线的宽度
    total_width = sum(img.width for img in processed_images) + SEPARATOR_WIDTH * (len(processed_images) - 1)
    stitched_image = Image.new('RGBA', (total_width, std_height), (255, 255, 255, 255))
    draw_stitched = ImageDraw.Draw(stitched_image)
    
    current_x = 0
    for i, img in enumerate(processed_images):
        stitched_image.paste(img, (current_x, 0))
        current_x += img.width
        
        # 如果不是最后一张图片，就在后面画一条分割线
        if i < len(processed_images) - 1:
            line_x = current_x + SEPARATOR_WIDTH // 2
            draw_stitched.line([(line_x, 0), (line_x, std_height)], fill=SEPARATOR_COLOR, width=SEPARATOR_WIDTH)
            current_x += SEPARATOR_WIDTH

    stitched_image.convert('RGB').save(output_path, 'JPEG', quality=90)


# --- 主执行区 ---
def main():
    # --- 请在这里修改你的路径 ---
    JSON_FILE_PATH = 'outputs/training_data_embedding/similarity_results_batched.json'
    IMAGE_DIRECTORY = ['/home/tfj/datasets/image_retri10k/images', '/home/tfj/datasets/image_retri10k/images_sup']
    BASE_OUTPUT_DIR = 'outputs/visual_output'
    TOP_N = 10
    
    os.makedirs(BASE_OUTPUT_DIR, exist_ok=True)
    print(f"结果将保存在: {os.path.abspath(BASE_OUTPUT_DIR)}")

    print(f"正在从 '{JSON_FILE_PATH}' 加载相似度数据...")
    try:
        with open(JSON_FILE_PATH, 'r', encoding='utf-8') as f:
            all_similarity_data = json.load(f)
    except (json.JSONDecodeError, FileNotFoundError) as e:
        print(f"错误：无法加载或解析JSON文件。 {e}")
        return
    
    all_query_filenames = list(all_similarity_data.keys())
    print(f"加载完成，共找到 {len(all_query_filenames)} 条记录。")
    time.sleep(1)

    pbar = tqdm(all_query_filenames, desc="生成可视化结果")
    for query_filename in pbar:
        try:
            item_id = get_item_id_from_filename(query_filename)
            output_subdir = os.path.join(BASE_OUTPUT_DIR, item_id)
            os.makedirs(output_subdir, exist_ok=True)
            
            output_filename = f"{os.path.splitext(query_filename)[0]}.jpg"
            final_output_path = os.path.join(output_subdir, output_filename)
            
            if os.path.exists(final_output_path):
                continue

            pbar.set_description(f"处理中: {query_filename}")

            visualize_one_result(
                query_image_filename=query_filename,
                similarity_data_for_query=all_similarity_data[query_filename],
                image_dir=IMAGE_DIRECTORY,
                output_path=final_output_path,
                top_n=TOP_N
            )
        except Exception as e:
            print(f"\n处理 '{query_filename}' 时发生错误: {e}")
            continue

    print("\n--- 所有可视化任务完成！ ---")

if __name__ == "__main__":
    main()